import getpass
import os
from dotenv import load_dotenv; load_dotenv()

from typing import Annotated

from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

from langchain_openai import ChatOpenAI
from langchain_core.language_models import BaseChatModel

from langchain_core.tools import BaseTool
from IPython.display import Image, display

from langchain_core.runnables import RunnableConfig
from utils.my_callback import MyCustomHandlerTwo, MyCustomAsyncHandler

from langchain import hub
import json
from langchain_core.messages import ToolMessage
from typing import Literal

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]


def chatbot(state: State):
    message = llm_with_tools.invoke(state["messages"])
    return {"messages": [message]}

def search_tool(debug=False) -> BaseTool :
    from langchain_community.tools.tavily_search import TavilySearchResults

    tool = TavilySearchResults(max_results=2)
    if debug:
        tool.invoke("What's a 'node' in LangGraph?")
    return tool

class BasicToolNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        for tool_call in message.tool_calls:
            if "properties" in tool_call["args"]:
                # qwen模型function_call
                tool_call_args = tool_call["args"]["properties"]
            else:
                tool_call_args = tool_call["args"]
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                tool_call_args
            )
            outputs.append(
                ToolMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"],
                    tool_call_id=tool_call["id"],
                )
            )
        return {"messages": outputs}

def route_tools(
    state: State,
) -> Literal["tools", "__end__"]:
    """
    Use in the conditional_edge to route to the ToolNode if the last message
    has tool calls. Otherwise, route to the end.
    """
    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
        return "tools"
    return "__end__"

# Optional, add tracing in LangSmith
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Multi-agent Collaboration"

_set_env("BAICHUAN_API_KEY")

# llm = ChatAnthropic(model="claude-3-haiku-20240307")
llm = ChatOpenAI(model="gpt-4o-mini")
# 本地百川模型
# llm = ChatOpenAI(model="Baichuan2", api_key="123", base_url="http://bc.192.168.107.2.nip.io/v1")
# 豆包大模型
# llm = ChatOpenAI(model="ep-20240801230159-vp5jn", api_key="eb64ed3f-f257-46d1-a046-c28cb553ca28", base_url="https://ark.cn-beijing.volces.com/api/v3")
# 本地千问模型, 支持function_call
# llm = ChatOpenAI(model="qwen2:72b", api_key="ollama", base_url="http://generate.yangzhiqiang.tech/v1")
# qwen模型
# llm = ChatOpenAI(model="qwen-plus", api_key="sk-5c7a7ef8061047cf98ba253b2ae06b2b", base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")
tool = search_tool()
llm_with_tools = llm.bind_tools([tool])

# Get the prompt to use - you can modify this!
prompt = hub.pull("hwchase17/openai-functions-agent")
print(f"prompt: {prompt.pretty_print()}")
# chat_agent = prompt | llm_with_tools

if __name__ == "__main__":

    handler2 = MyCustomHandlerTwo()
    handler3 = MyCustomAsyncHandler()

    graph_builder = StateGraph(State)

    tool_node = BasicToolNode(tools=[tool])
    graph_builder.add_node("tools", tool_node)

    # The first argument is the unique node name
    # The second argument is the function or object that will be called whenever
    # the node is used.
    graph_builder.add_node("chatbot", chatbot)

    graph_builder.add_edge(START, "chatbot")
    # The `tools_condition` function returns "tools" if the chatbot asks to use a tool, and "__end__" if
    # it is fine directly responding. This conditional routing defines the main agent loop.
    graph_builder.add_conditional_edges(
        "chatbot",
        route_tools,
        # The following dictionary lets you tell the graph to interpret the condition's outputs as a specific node
        # It defaults to the identity function, but if you
        # want to use a node named something else apart from "tools",
        # You can update the value of the dictionary to something else
        # e.g., "tools": "my_tools"
        {"tools": "tools", "__end__": "__end__"},
    )
    # Any time a tool is called, we return to the chatbot to decide the next step
    graph_builder.add_edge("tools", "chatbot")
    graph = graph_builder.compile()

    try:
        # 提取文件名前缀，不包含目录
        graph.get_graph().draw_mermaid_png(output_file_path=os.path.dirname(__file__) + "/" + os.path.basename(__file__).split(".")[0] + "_graph.png")
        # display(Image(graph.get_graph().draw_mermaid_png()))
    except Exception:
        # This requires some extra dependencies and is optional
        print(f"Error: {Exception}")
        pass

    config = RunnableConfig(callbacks=[handler3])
    while True:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        for event in graph.stream({"messages": ("user", user_input)}, config=config):
            for value in event.values():
                print("Assistant:", value["messages"][-1].content)

